package main

import (
	"compress/gzip"
	"log"
	"net/url"
	"os"
	"os/signal"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	// 设置中断信号通道
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	// 定义 WebSocket 服务器地址
	u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/demo"}
	log.Printf("connecting to %s", u.String())

	// 建立 WebSocket 连接
	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer func(c *websocket.Conn) {
		_ = c.Close()
	}(c)

	// 完成通道用于通知主 goroutine 工作完成
	done := make(chan struct{})

	// 接收消息的 goroutine
	go func() {
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				return
			}
			log.Printf("recv: %s", message)
		}
	}()

	// 发送消息的 goroutine
	go func() {
		ticker := time.NewTicker(time.Minute * 1)
		pingTicker := time.NewTicker(9 * time.Second)
		defer ticker.Stop()
		defer pingTicker.Stop()

		for {
			select {
			case <-done:
				return
			case t := <-ticker.C:
				writeErr := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
				if writeErr != nil {
					log.Println("write text message:", writeErr)
					return
				}

				writeErr = writeGzipMessage(c, []byte(t.String()))
				if writeErr != nil {
					log.Println("write gzip message:", writeErr)
					return
				}

			case <-pingTicker.C:
				// 发送 ping 消息
				writeErr := c.WriteMessage(websocket.PingMessage, nil)
				if err != nil {
					log.Println("write ping:", writeErr)
					return
				}
			}

		}
	}()

	// 主循环，等待中断信号
	for {
		select {
		case <-done:
			return
		case <-interrupt:
			log.Println("interrupt")

			// 发送关闭消息通知服务器
			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Println("write close:", err)
				return
			}

			// 等待接收 goroutine 完成
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			return
		}
	}
}

func writeGzipMessage(conn *websocket.Conn, msg []byte) error {
	var writeWait = 10 * time.Second
	_ = conn.SetWriteDeadline(time.Now().Add(writeWait))
	wsWriter, err := conn.NextWriter(websocket.BinaryMessage)
	if err != nil {
		return err
	}

	gzWriter, err := gzip.NewWriterLevel(wsWriter, 3)
	if err != nil {
		return err
	}
	if _, err := gzWriter.Write(msg); err != nil {
		return err
	}

	if err := gzWriter.Close(); err != nil {
		return err
	}

	return nil
}
